#include <windows.h>
#include <mfapi.h>
#include <mfidl.h>
#include <mfreadwrite.h>
#include <mfplay.h>

#include "camera.h"

/**
source --> callback

source --> player
       |_ callback

*/


#define SAFE_RELEASE(comobj) if(comobj)comobj->lpVtbl->Release(comobj)



static int init(){
    static int inited = 0;
    if(inited)return 0;
    CoInitialize(NULL);
    MFStartup(MF_VERSION, MFSTARTUP_LITE);    
    inited = 1;
    return 0;
}


struct _camera{
    const IMFSourceReaderCallbackVtbl* lpVtbl;
    ULONG ref;
    IMFMediaSource* source;
    IMFSourceReader* reader;
    //IMFMediaType* type;
    camera_data_callback on_data;
    void* udata;
    uint32_t type;
    int width;
    int height;
};

static HRESULT __stdcall Camera_QueryInterface(IMFSourceReaderCallback* pThis, REFIID iid, void** arg){
    if (IsEqualGUID(iid, &IID_IMFSourceReaderCallback)){
        *arg = pThis;
        return S_OK;
    }
    return E_NOINTERFACE;
}
    

static ULONG __stdcall Camera_AddRef(IMFSourceReaderCallback* pThis){
    camera_t c = (camera_t)pThis;
    c->ref ++;
    return c->ref;
}

static ULONG __stdcall Camera_Release(IMFSourceReaderCallback* pThis){
    camera_t c = (camera_t)pThis;
    c->ref --;
    if(c->ref <= 0){
        free(c);
        return 0;
    }
    return c->ref;
}

static HRESULT __stdcall Camera_OnEvent(IMFSourceReaderCallback* pThis, DWORD dwStreamIndex, IMFMediaEvent* pEvent){
    return S_OK;
}

static HRESULT __stdcall Camera_OnFlush(IMFSourceReaderCallback* pThis, DWORD dwStreamIndex){
    return S_OK;
}

static HRESULT __stdcall Camera_OnReadSample(IMFSourceReaderCallback* pThis, HRESULT hrStatus, DWORD dwStreamIndex, DWORD dwStreamFlags, LONGLONG llTimestamp, IMFSample *pSample){
    HRESULT hr;
    camera_t c = (camera_t)pThis;
    IMFMediaBuffer *pBuffer = NULL;
    BYTE* buf = NULL;
    DWORD dwLen = 0;
    if(FAILED(hrStatus))return hrStatus;
    if(pSample){        
        hr = pSample->lpVtbl->GetBufferByIndex(pSample, 0, &pBuffer);
        if(!FAILED(hr)){
            hr = pBuffer->lpVtbl->Lock(pBuffer, &buf, NULL, &dwLen);
            if(!FAILED(hr)){
                c->on_data(buf, dwLen, c->udata);
                pBuffer->lpVtbl->Unlock(pBuffer);
            }
            SAFE_RELEASE(pBuffer);
        }
    }
    hr = c->reader->lpVtbl->ReadSample(c->reader, (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM,0,NULL,NULL,NULL,NULL);
    if(FAILED(hr))return hr;       
    return S_OK;
}

static const IMFSourceReaderCallbackVtbl IMFSourceReaderCallbackVtbl_ = {
    Camera_QueryInterface,
    Camera_AddRef,
    Camera_Release,
    Camera_OnReadSample,
    Camera_OnFlush,    
    Camera_OnEvent,
};


int camera_list(camera_list_callback cb, void* udata){
    HRESULT hr;
    IMFAttributes* pAttributes = NULL;
    IMFActivate** ppDevices = NULL; 
    UINT32 i, count, len;
    LPWSTR wname = NULL;
    char name[256];
    int ret = -1;
    count = 0;  
    if(init())return -1;
    hr = MFCreateAttributes(&pAttributes, 1);
    if(FAILED(hr))return -1;
    hr = pAttributes->lpVtbl->SetGUID(pAttributes, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID);
    if(FAILED(hr))goto clean;
    hr = MFEnumDeviceSources(pAttributes, &ppDevices, &count);
    if(FAILED(hr))goto clean;
    for(i = 0; i < count; i++){
        hr = ppDevices[i]->lpVtbl->GetAllocatedString(ppDevices[i], &MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &wname, &len);
        if(FAILED(hr))goto clean;
        name[0] = 0;
        WideCharToMultiByte(CP_UTF8, 0, wname, -1, name, 256, NULL, NULL);
        cb(name, udata);
        CoTaskMemFree(wname);
        wname = NULL;
    }
    ret = 0;   
clean:
    for(i = 0; i < count; i++){
        SAFE_RELEASE(ppDevices[i]);
    }    
    SAFE_RELEASE(pAttributes);
    if(ppDevices)CoTaskMemFree(ppDevices);
    if(wname)CoTaskMemFree(wname);
    
    return ret;
}

#define CHECK_TYPE(type, fcc) if(IsEqualGUID(&guid, &type))return fourccs(fcc)

static uint32_t get_type(IMFMediaType* pType){
    GUID guid;
    pType->lpVtbl->GetGUID(pType, &MF_MT_SUBTYPE, &guid);
    CHECK_TYPE(MFVideoFormat_YUY2,"YUY2"); //     FCC('YUY2')
    //CHECK_TYPE(MFVideoFormat_RGB24,"RGB "); //    D3DFMT_R8G8B8 
    //CHECK_TYPE(MFVideoFormat_RGB32,"RGBA"); //    D3DFMT_X8R8G8B8 
    CHECK_TYPE(MFVideoFormat_NV11,"NV11"); //     FCC('NV11')
    CHECK_TYPE(MFVideoFormat_NV12,"NV12"); //     FCC('NV12') 
    CHECK_TYPE(MFVideoFormat_MJPG,"MJPG");          
    CHECK_TYPE(MFVideoFormat_AI44,"AI44"); //     FCC('AI44')
    //CHECK_TYPE(MFVideoFormat_ARGB32,"ARGB"); //   D3DFMT_A8R8G8B8 
    CHECK_TYPE(MFVideoFormat_AYUV,"AYUV"); //     FCC('AYUV')
    CHECK_TYPE(MFVideoFormat_DV25,"DV25"); //     FCC('dv25')
    CHECK_TYPE(MFVideoFormat_DV50,"DV50"); //     FCC('dv50')
    CHECK_TYPE(MFVideoFormat_DVH1,"DVH1"); //     FCC('dvh1')
    CHECK_TYPE(MFVideoFormat_DVSD,"DVSD"); //     FCC('dvsd')
    CHECK_TYPE(MFVideoFormat_DVSL,"DVSL"); //     FCC('dvsl')
    CHECK_TYPE(MFVideoFormat_H264,"H264"); //     FCC('H264')
    CHECK_TYPE(MFVideoFormat_I420,"I420"); //     FCC('I420')
    CHECK_TYPE(MFVideoFormat_IYUV,"IYUV"); //     FCC('IYUV')
    CHECK_TYPE(MFVideoFormat_M4S2,"M4S2"); //     FCC('M4S2')
    CHECK_TYPE(MFVideoFormat_MP43,"MP43"); //     FCC('MP43')
    CHECK_TYPE(MFVideoFormat_MP4S,"MP4S"); //     FCC('MP4S')
    CHECK_TYPE(MFVideoFormat_MP4V,"MP4V"); //     FCC('MP4V')
    CHECK_TYPE(MFVideoFormat_MPG1,"MPG1"); //     FCC('MPG1')
    CHECK_TYPE(MFVideoFormat_MSS1,"MSS1"); //     FCC('MSS1')
    CHECK_TYPE(MFVideoFormat_MSS2,"MSS2"); //     FCC('MSS2')

    CHECK_TYPE(MFVideoFormat_P010,"P010"); //     FCC('P010')
    CHECK_TYPE(MFVideoFormat_P016,"P016"); //     FCC('P016')
    CHECK_TYPE(MFVideoFormat_P210,"P210"); //     FCC('P210')
    CHECK_TYPE(MFVideoFormat_P216,"P216"); //     FCC('P216')

    //CHECK_TYPE(MFVideoFormat_RGB555,"RGB555"); //   D3DFMT_X1R5G5B5 
    //CHECK_TYPE(MFVideoFormat_RGB565,"RGB565"); //   D3DFMT_R5G6B5 
    CHECK_TYPE(MFVideoFormat_RGB8,"RGB8");
    CHECK_TYPE(MFVideoFormat_UYVY,"UYVY"); //     FCC('UYVY')
    CHECK_TYPE(MFVideoFormat_v210,"v210"); //     FCC('v210')
    CHECK_TYPE(MFVideoFormat_v410,"v410"); //     FCC('v410')
    CHECK_TYPE(MFVideoFormat_WMV1,"WMV1"); //     FCC('WMV1')
    CHECK_TYPE(MFVideoFormat_WMV2,"WMV2"); //     FCC('WMV2')
    CHECK_TYPE(MFVideoFormat_WMV3,"WMV3"); //     FCC('WMV3')
    CHECK_TYPE(MFVideoFormat_WVC1,"WVC1"); //     FCC('WVC1')
    CHECK_TYPE(MFVideoFormat_Y210,"Y210"); //     FCC('Y210')
    CHECK_TYPE(MFVideoFormat_Y216,"Y216"); //     FCC('Y216')
    CHECK_TYPE(MFVideoFormat_Y410,"Y410"); //     FCC('Y410')
    CHECK_TYPE(MFVideoFormat_Y416,"Y416"); //     FCC('Y416')
    CHECK_TYPE(MFVideoFormat_Y41P,"Y41P");
    CHECK_TYPE(MFVideoFormat_Y41T,"Y41T");
    //CHECK_TYPE(MFVideoFormat_YUY2,"YUY2"); //     FCC('YUY2')
    CHECK_TYPE(MFVideoFormat_YV12,"YV12"); //     FCC('YV12')
    CHECK_TYPE(MFVideoFormat_YVYU,"YVYU");
    return 0;
}

static int get_size(IMFMediaType* pType, int* w, int* h){
    UINT64 sz;
    HRESULT hr = pType->lpVtbl->GetUINT64(pType, &MF_MT_FRAME_SIZE, &sz);    
    if(FAILED(hr))return -1;
    *w = sz >> 32;
    *h = sz & 0xffffffff;  
    return 0;
}


/*
static void LogMediaType(IMFMediaType* pType){
    uint32_t type = 0;
    int w = 0;
    int h = 0;
    char name[6] = {0};
    type = get_type(pType);
    get_size(pType,&w,&h);
    memcpy(name, &type, 4);
    printf("%d %d %s\n", w, h, name);
}


HRESULT EnumerateCaptureFormats(IMFMediaSource *pSource)
{
    IMFPresentationDescriptor *pPD = NULL;
    IMFStreamDescriptor *pSD = NULL;
    IMFMediaTypeHandler *pHandler = NULL;
    IMFMediaType *pType = NULL;

    HRESULT hr = pSource->lpVtbl->CreatePresentationDescriptor(pSource, &pPD);
    if (FAILED(hr))
    {
        goto done;
    }

    BOOL fSelected;
    hr = pPD->lpVtbl->GetStreamDescriptorByIndex(pPD, 0, &fSelected, &pSD);
    if (FAILED(hr))
    {
        goto done;
    }

    hr = pSD->lpVtbl->GetMediaTypeHandler(pSD,&pHandler);
    if (FAILED(hr))
    {
        goto done;
    }

    DWORD cTypes = 0;
    hr = pHandler->lpVtbl->GetMediaTypeCount(pHandler,&cTypes);
    if (FAILED(hr))
    {
        goto done;
    }

    for (DWORD i = 0; i < cTypes; i++)
    {
        hr = pHandler->lpVtbl->GetMediaTypeByIndex(pHandler,i, &pType);
        if (FAILED(hr))
        {
            goto done;
        }

        LogMediaType(pType);
        //OutputDebugString(L"\n");

        SAFE_RELEASE(pType);
    }

done:
    SAFE_RELEASE(pPD);
    SAFE_RELEASE(pSD);
    SAFE_RELEASE(pHandler);
    SAFE_RELEASE(pType);
    return hr;
}
*/

static int set_type_by_flag(IMFMediaSource* pSource, int flag){
    int ret = -1;
    IMFPresentationDescriptor *pPD = NULL;
    IMFStreamDescriptor *pSD = NULL;
    IMFMediaTypeHandler *pHandler = NULL;
    IMFMediaType* pType = NULL;
    HRESULT hr;
    BOOL fSelected;
    DWORD cTypes = 0;
    IMFMediaType* pTypeSel = NULL;
    DWORD i;
    int width = 0;
    int height = 0;    
    GUID guid = {0};
    int w,h;
    int selindex = -1;    
    hr = pSource->lpVtbl->CreatePresentationDescriptor(pSource, &pPD);
    if (FAILED(hr))goto clean; 
    hr = pPD->lpVtbl->GetStreamDescriptorByIndex(pPD, 0, &fSelected, &pSD);
    if (FAILED(hr))goto clean;
    hr = pSD->lpVtbl->GetMediaTypeHandler(pSD,&pHandler);
    if (FAILED(hr))goto clean;   
    hr = pHandler->lpVtbl->GetMediaTypeCount(pHandler,&cTypes);
    if (FAILED(hr))goto clean;
    hr = pHandler->lpVtbl->GetCurrentMediaType(pHandler, &pTypeSel);
    if (FAILED(hr))goto clean; 

    if(flag & 0x0F){ //需要设置大小
        if(flag & CAMERA_MIN_RESOLUTION){
            width = 99999; height = 99999;
        }else if(flag & CAMERA_MAX_RESOLUTION){
            width = 0; height = 0;
        }
        for (i = 0; i < cTypes; i++){
            hr = pHandler->lpVtbl->GetMediaTypeByIndex(pHandler,i, &pType);
            if (FAILED(hr))goto clean;        
            get_size(pType,&w,&h);
            if(flag & CAMERA_MIN_RESOLUTION){
                if(w < width && h < height){
                    width = w;
                    height = h;
                    selindex = i;
                }
            }else if(flag & CAMERA_MAX_RESOLUTION){
                if(w > width && h > height){
                    width = w;
                    height = h;
                    selindex = i;                
                }              
            }
            SAFE_RELEASE(pType);
            pType = NULL;
        }
    }else{
        get_size(pTypeSel,&width,&height);
    }

    if(flag & 0xF0){ //需要设置像素格式
        for (i = 0; i < cTypes; i++){
            hr = pHandler->lpVtbl->GetMediaTypeByIndex(pHandler,i, &pType);
            if (FAILED(hr))goto clean;
            get_size(pType,&w,&h); 
            if(w != width || h != height){
                SAFE_RELEASE(pType);
                pType = NULL;
                continue;
            }
            pType->lpVtbl->GetGUID(pType, &MF_MT_SUBTYPE, &guid);
            if(flag & CAMERA_PREFER_RGB){                
                pType->lpVtbl->GetGUID(pType, &MF_MT_SUBTYPE, &guid);
                if(IsEqualGUID(&guid, &MFVideoFormat_RGB24) ||
                    IsEqualGUID(&guid, &MFVideoFormat_RGB32) ||
                    IsEqualGUID(&guid, &MFVideoFormat_RGB555) ||
                    IsEqualGUID(&guid, &MFVideoFormat_RGB565) ||
                    IsEqualGUID(&guid, &MFVideoFormat_ARGB32) ){                        
                        selindex = i;     
                        break;                       
                }
            }else if(flag & CAMERA_PREFER_GRAY){
                pType->lpVtbl->GetGUID(pType, &MF_MT_SUBTYPE, &guid);
                if(IsEqualGUID(&guid, &MFVideoFormat_NV11) ||
                    IsEqualGUID(&guid, &MFVideoFormat_NV12) ||
                    //IsEqualGUID(&guid, &MFVideoFormat_NV21) ||
                    IsEqualGUID(&guid, &MFVideoFormat_YV12) ){
                        selindex = i; 
                        break;                  
                } 
            }
            SAFE_RELEASE(pType);  
            pType = NULL;
        }   
    }

    if(selindex >= 0){
        hr = pHandler->lpVtbl->GetMediaTypeByIndex(pHandler,selindex, &pType);
        if (FAILED(hr))goto clean;        
        hr = pHandler->lpVtbl->SetCurrentMediaType(pHandler, pType);
        if (FAILED(hr))goto clean;
    }
    ret = 0;   
clean:
    SAFE_RELEASE(pPD);
    SAFE_RELEASE(pSD);
    SAFE_RELEASE(pHandler);
    SAFE_RELEASE(pType);
    SAFE_RELEASE(pTypeSel);
    return ret;
}


camera_t camera_open(int index, int flag){
    HRESULT hr;
    IMFMediaSource* pSource = NULL;
    IMFSourceReader* pReader = NULL;
    IMFAttributes* pAttributes = NULL;
    IMFAttributes* pReaderAttributes = NULL;
    IMFActivate** ppDevices = NULL; 
    IMFMediaType* pType = NULL;
    UINT32 i,count;
    camera_t c;
    int ret = -1;    
    count = 0;
    if(init())return NULL;
    c = (camera_t)malloc(sizeof(struct _camera));
    if(!c)return NULL;
    hr = MFCreateAttributes(&pAttributes, 1);
    if(FAILED(hr))goto clean;
    hr = pAttributes->lpVtbl->SetGUID(pAttributes, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, &MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID);
    if(FAILED(hr))goto clean;
    hr = MFEnumDeviceSources(pAttributes, &ppDevices, &count);
    if(FAILED(hr))goto clean;
    if(index < 0 || index >= count)goto clean;
    hr = ppDevices[index]->lpVtbl->ActivateObject(ppDevices[index], &IID_IMFMediaSource, (LPVOID*)&pSource);
    if(FAILED(hr))goto clean;
    if(flag)set_type_by_flag(pSource, flag);
    //EnumerateCaptureFormats(pSource);
    hr = MFCreateAttributes(&pReaderAttributes, 2);
    if(FAILED(hr))goto clean;
    hr = pReaderAttributes->lpVtbl->SetUINT32(pReaderAttributes, &MF_READWRITE_DISABLE_CONVERTERS, TRUE);
    if(FAILED(hr))goto clean;
    c->lpVtbl = &IMFSourceReaderCallbackVtbl_;
    c->ref = 1;
    c->source = pSource;
    c->reader = NULL;
    c->on_data = NULL;
    c->udata = NULL;
    hr = pReaderAttributes->lpVtbl->SetUnknown(pReaderAttributes, &MF_SOURCE_READER_ASYNC_CALLBACK, (IUnknown*)c);
    if(FAILED(hr))goto clean;   
    hr = MFCreateSourceReaderFromMediaSource(pSource, pReaderAttributes, &pReader);
    if(FAILED(hr))goto clean;
    hr = pReader->lpVtbl->GetCurrentMediaType(pReader, (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM, &pType);
    if(FAILED(hr))goto clean;
    c->type = get_type(pType);
    if(!c->type)goto clean;
    if(get_size(pType, &c->width, &c->height))goto clean;
    hr = pReader->lpVtbl->ReadSample(pReader, (DWORD)MF_SOURCE_READER_FIRST_VIDEO_STREAM,0,NULL,NULL,NULL,NULL);
    if(FAILED(hr))goto clean;
    c->reader = pReader;
    ret = 0; 
clean:
    for(i = 0; i < count; i++){
        SAFE_RELEASE(ppDevices[i]);
    }    
    SAFE_RELEASE(pReaderAttributes);
    SAFE_RELEASE(pAttributes);
    SAFE_RELEASE(pType);
    if(ppDevices)CoTaskMemFree(ppDevices); 
    if(ret){
        SAFE_RELEASE(pReader);
        SAFE_RELEASE(pSource);
        free(c);
        c = NULL;
    }  
    return c;
}

int camera_set_callback(camera_t c, camera_data_callback cb, void* udata){
    c->on_data = cb;
    c->udata = udata;
    return 0;
}

uint32_t camera_get_type(camera_t c){
    return c->type;
}

int camera_get_width(camera_t c){
    return c->width;
}

int camera_get_height(camera_t c){
    return c->height;
}

